# -*- coding: utf-8 -*-
# time: 2025/5/10 09:18
# file: tf01.py
# author: hanson
"""
小模型推理
Pythia-70m需要约1GB GPU内存

"""
from datasets import load_dataset



# 保存微调数据集（可选）
# fine_tuning_dataset.save_to_disk("./modified_lamini_docs")  # 保存到本地


def get_modify_dataSet(a: int) -> dict:
    # 加载微调数据集
    fine_tuning_dataset_file = "lamini/lamini_docs"
    fine_tuning_dataset = load_dataset(fine_tuning_dataset_file)
    print(fine_tuning_dataset)

    # 查看文件数据，并在 train和test数增加一行数据
    print(fine_tuning_dataset["train"][0])  # 查看第一条训练数据
    print(fine_tuning_dataset["test"][0])  # 查看第一条测试数据

    # 根据实际列名补充
    new_train_row = {'question': '中国中山大学在哪里?', 'answer': "中国中山大学在哪里？", 'input_ids': [],
                     'attention_mask': [], 'labels': []}
    new_train_row2 = {'question': '请用中文介绍自己？',
                      'answer': '我叫laich,来自中国广东省中山市, 由Hanson开发的人工智能对话机器人。我很智能的，不信你可以提问中国中山大学在哪里？我会告诉你的。',
                      'input_ids': [],
                      'attention_mask': [], 'labels': []}

    new_test_row = {'question': '中国中山大学在哪里?', 'answer': "中国中山大学在哪里？", 'input_ids': [],
                    'attention_mask': [], 'labels': []}
    new_test_row2 = {'question': '请用中文介绍自己？',
                     'answer': '我叫laich,来自中国广东省中山市, 由Hanson开发的人工智能对话机器人。我很智能的，不信你可以提问中国中山大学在哪里？我会告诉你的。',
                     'input_ids': [],
                     'attention_mask': [],
                     'labels': []}
    # 添加训练集新行
    fine_tuning_dataset["train"] = fine_tuning_dataset["train"].add_item(new_train_row)
    fine_tuning_dataset["train"] = fine_tuning_dataset["train"].add_item(new_train_row2)
    # 添加测试集新行
    fine_tuning_dataset["test"] = fine_tuning_dataset["test"].add_item(new_test_row)
    fine_tuning_dataset["test"] = fine_tuning_dataset["test"].add_item(new_test_row2)
    # 验证添加结果
    print("Last train sample:", fine_tuning_dataset["train"][-1])  # 查看最后一条训练数据
    print("Last test sample:", fine_tuning_dataset["test"][-1])  # 查看最后一条测试数据
    if (1 == a):
        return {"dataSet": fine_tuning_dataset}

    return {"dataSet": load_dataset(fine_tuning_dataset_file)}
